package adkr

import (
    "encoding/json"
    "log"
    "sync"
)

// Strong Provable Broadcast (SPBC) simplified interface and implementation.
// This follows the classic VAL/ECHO/READY pattern with a proof (here: commit+merkle root) as payload.

type SPBCTag string

const (
    SPBCValTag   SPBCTag = "VAL"
    SPBCEchoTag  SPBCTag = "ECHO"
    SPBCReadyTag SPBCTag = "READY"
)

type SPBCMessage struct {
    Tag     SPBCTag `json:"tag"`
    Round   int     `json:"round"`
    Dealer  int     `json:"dealer"`
    Payload []byte  `json:"payload"`
}

// SPBCVal carries the commitment root and the encoded stripe for the receiver
type SPBCVal struct {
    RootHash []byte   `json:"root_hash"`
    Branch   [][]byte `json:"branch"`
    Stripe   []byte   `json:"stripe"`
}

type spbcState struct {
    mu          sync.Mutex
    n, f        int
    round       int
    dealer      int
    rootHash    []byte
    echoSet     map[int]struct{}
    readySet    map[int]struct{}
    delivered   bool
}

// StrongProvableBroadcast runs a single SPBC instance for the given dealer.
// - sid: unique session id
// - pid: local process id
// - n,f: system params, n>=3f+1
// - dealer: dealer process id
// - inputVal: only used by dealer to provide the VAL payload (SPBCVal JSON bytes)
// - msgInCh/msgOutCh: network channels
// - deliver: callback invoked with reconstructed payload (root + stripe) when delivered
func StrongProvableBroadcast(
    sid string,
    pid int,
    n, f int,
    dealer int,
    round int,
    inputVal []byte,
    msgInCh <-chan Message,
    msgOutCh chan<- Message,
    deliver func(val SPBCVal),
) {
    st := &spbcState{
        n:       n,
        f:       f,
        round:   round,
        dealer:  dealer,
        echoSet: make(map[int]struct{}),
        readySet: make(map[int]struct{}),
    }

    // Dealer sends VAL
    if pid == dealer && inputVal != nil {
        // inputVal is already an SPBCVal marshaled as JSON (root/branch/stripe)
        spbcMsg := SPBCMessage{Tag: SPBCValTag, Round: round, Dealer: dealer, Payload: inputVal}
        payload, _ := json.Marshal(spbcMsg)
        // broadcast
        msgOutCh <- Message{Tag: BroadcastTag(ADKR_SPBC), Round: round, To: -1, Payload: payload}
        log.Printf("SPBC-SID:%s PID:%d [Dealer %d] broadcast VAL", sid, pid, dealer)
    }

    for {
        msg, ok := <-msgInCh
        if !ok { return }
        if msg.Tag != BroadcastTag(ADKR_SPBC) || msg.Round != round { continue }

        var spbcMsg SPBCMessage
        if err := json.Unmarshal(msg.Payload, &spbcMsg); err != nil { continue }
        if spbcMsg.Dealer != dealer { continue }

        switch spbcMsg.Tag {
        case SPBCValTag:
            // upon VAL from dealer, send ECHO with same payload
            st.mu.Lock()
            var val SPBCVal
            if err := json.Unmarshal(spbcMsg.Payload, &val); err == nil {
                st.rootHash = val.RootHash
            }
            st.mu.Unlock()

            echo := SPBCMessage{Tag: SPBCEchoTag, Round: round, Dealer: dealer, Payload: spbcMsg.Payload}
            payload, _ := json.Marshal(echo)
            msgOutCh <- Message{Tag: BroadcastTag(ADKR_SPBC), Round: round, To: -1, Payload: payload}

        case SPBCEchoTag:
            st.mu.Lock()
            st.echoSet[msg.Sender] = struct{}{}
            eCnt := len(st.echoSet)
            root := st.rootHash
            st.mu.Unlock()
            // when received at least n - f ECHO, send READY
            if eCnt >= st.n - st.f && len(root) > 0 {
                ready := SPBCMessage{Tag: SPBCReadyTag, Round: round, Dealer: dealer, Payload: spbcMsg.Payload}
                payload, _ := json.Marshal(ready)
                msgOutCh <- Message{Tag: BroadcastTag(ADKR_SPBC), Round: round, To: -1, Payload: payload}
            }

        case SPBCReadyTag:
            st.mu.Lock()
            st.readySet[msg.Sender] = struct{}{}
            rCnt := len(st.readySet)
            delivered := st.delivered
            st.mu.Unlock()
            // deliver when received f+1 READY (or 2f+1 depending on variant); using f+1 here to unblock pipeline
            if !delivered && rCnt >= st.f+1 {
                var val SPBCVal
                if err := json.Unmarshal(spbcMsg.Payload, &val); err == nil {
                    st.mu.Lock(); st.delivered = true; st.mu.Unlock()
                    // Callback deliver
                    deliver(val)
                    return
                }
            }
        }
    }
}


